import matplotlib.pyplot as plt
import numpy as np


if __name__ == "__main__":
    q10 = 2.3
    celsius = 25
    shift = 2.96
    qt = q10 ** ((celsius - 22) / 10)

    def tau_g(v):
        tau_n = (0.5 + 0.124 * np.exp(-v / 15.8)) / qt
        tau_l = (10.4 + 0.0118 * np.exp(-v / 7.85)) / qt
        return tau_n, tau_l

    def tau_h(v):

        tau_n = (0.774 + 0.14 * np.exp(-v / 13.27)) / qt
        tau_l = (22.25 + 0.0455 * np.exp(-v / 7.46)) / qt
        return tau_n, tau_l

    def tau_i(v):
        tau_n = (2.71 + 0.028 * np.exp(-v / 9.34)) / qt
        tau_l = (110 + 0.0009 * np.exp(-v / 5.16)) / qt
        return tau_n, tau_l

    def tau_it2(v):

        phi_m = 5.0 ** ((celsius - 24) / 10)
        phi_h = 3.0 ** ((celsius - 24) / 10)
        tau_m = (
            3 + 1.0 / (np.exp((v + shift + 25) / 10) + np.exp(-(v + shift + 100) / 15))
        ) / phi_m
        tau_h = (
            85 + 1.0 / (np.exp((v + shift + 46) / 4) + np.exp(-(v + shift + 405) / 50))
        ) / phi_h
        return tau_m, tau_h

    def tau_it2_fast(v):

        phi_m = 5.0 ** ((celsius - 24) / 10)
        phi_h = 3.0 ** ((celsius - 24) / 10)
        tau_m = (
            0.5 + 1.0 / (np.exp((v + shift + 25) / 15) + np.exp(-(v + shift + 100) / 15))
        ) / phi_m
        tau_h = (
            10 + 1.0 / (np.exp((v + shift + 46) / 5) + np.exp(-(v + shift + 405) / 50))
        ) / phi_h
        return tau_m, tau_h

    vs = np.linspace(-100, -20, 100)
    fig, ax = plt.subplots()
    plt.plot(vs, tau_g(vs)[0], label="Cav 3.1 (g)", c="C0")
    plt.plot(vs, tau_h(vs)[0], label="Cav3.2 (h)", c="C1")
    plt.plot(vs, tau_i(vs)[0], label="Cav3.3 (i)", c="C2")
    plt.plot(vs, tau_it2(vs)[0], label="it2", c="C3")
    plt.plot(vs, tau_it2_fast(vs)[0], label="it2_fast", c="C4")
    plt.axis([-100, -20, 0, 50])
    plt.legend()
    plt.savefig("tau_act.pdf")

    fig, ax = plt.subplots()
    plt.plot(vs, tau_g(vs)[1], label="Cav 3.1 (g)", c="C0", ls="--")
    plt.plot(vs, tau_h(vs)[1], label="Cav3.2 (h)", c="C1", ls="--")
    plt.plot(vs, tau_i(vs)[1], label="Cav3.3 (i)", c="C2", ls="--")
    plt.plot(vs, tau_it2(vs)[1], label="it2", c="C3", ls="--")
    plt.plot(vs, tau_it2_fast(vs)[1], label="it2_fast", c="C4", ls="--")
    plt.axis([-100, -20, 0, 700])
    plt.legend()
    plt.savefig("tau_inact.pdf")

    def inf_g(v):
        vhalfn = -51.73
        vhalfl = -85.4

        kn = 6.53
        kl = -5.4
        ninf = 1 / (1 + np.exp(-(v - vhalfn) / kn))
        linf = 1 / (1 + np.exp(-(v - vhalfl) / kl))
        return ninf, linf

    def inf_h(v):

        vhalfn = -43.15
        vhalfl = -73.9
        kn = 5.34
        kl = -2.76

        ninf = 1 / (1 + np.exp(-(v - vhalfn) / kn))
        linf = 1 / (1 + np.exp(-(v - vhalfl) / kl))
        return ninf, linf

    def inf_i(v):
        vhalfn = -60.7
        vhalfl = -93.2
        kn = 8.39
        kl = -5.4
        ninf = 1 / (1 + np.exp(-(v - vhalfn) / kn))
        linf = 1 / (1 + np.exp(-(v - vhalfl) / kl))
        return ninf, linf

    def inf_it2(v):
        m_inf = 1.0 / (1 + np.exp(-(v + shift + 50) / 7.4))
        h_inf = 1.0 / (1 + np.exp((v + shift + 78) / 5.0))
        return m_inf, h_inf

    vs = np.linspace(-100, -20, 100)
    fig, ax = plt.subplots(figsize=(6, 3))
    plt.plot(vs, inf_g(vs)[0], label="Cav 3.1 (g)", c="C0")
    plt.plot(vs, inf_g(vs)[1], c="C0", ls="--")
    plt.plot(vs, inf_h(vs)[0], label="Cav3.2 (h)", c="C1")
    plt.plot(vs, inf_h(vs)[1], c="C1", ls="--")
    plt.plot(vs, inf_i(vs)[0], label="Cav3.3 (i)", c="C2")
    plt.plot(vs, inf_i(vs)[1], c="C2", ls="--")
    plt.plot(vs, inf_it2(vs)[0], label="it2", c="C3")
    plt.plot(vs, inf_it2(vs)[1], c="C3", ls="--")
    plt.legend()
    plt.axis([-90, -50, 0, 0.2])
    plt.xlabel("v")
    plt.ylabel("inf")
    plt.tight_layout()
    plt.savefig("infs.pdf")
